package com.zhai.boot.interceptor;

import com.google.common.collect.ImmutableList;
import com.zhai.boot.annotation.LimitStream;
import com.zhai.boot.common.ServiceResult;
import com.zhai.boot.common.ServiceResultConstant;
import com.zhai.boot.util.RedisUtil;
import com.zhai.boot.util.StringUtil;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.http.ResponseEntity;
import org.springframework.scripting.support.ResourceScriptSource;

import java.lang.reflect.Method;

/**
 * 限流
 */
@Aspect
@Configuration
public class LimitStreamInterceptor {

    @Autowired
    private RedisUtil redisUtil;

    private static final String LIMIT_STREAM_KEY = "limitstream:";
    private static final String LIMIT_STREAM_LUA_PATH = "lua/limitstream.lua";

    @Around("execution(public * *(..)) && @annotation(com.zhai.boot.annotation.LimitStream)")
    public Object interceptor(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        LimitStream limitStream = method.getAnnotation(LimitStream.class);
        Integer limitPeriod = limitStream.period();
        Integer limitCount = limitStream.count();
        ImmutableList<String> keys = ImmutableList.of(StringUtil.join(LIMIT_STREAM_KEY,limitStream.prefix()));

        DefaultRedisScript<Number> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(LIMIT_STREAM_LUA_PATH)));
        redisScript.setResultType(Number.class);

        Number count = redisUtil.execute(redisScript, keys, limitCount, limitPeriod);
        if (count != null && count.intValue() <= limitCount) {
            return pjp.proceed();
        } else {
            return ResponseEntity.ok().body(ServiceResult.failure(ServiceResultConstant.SYS_LIMIT_STREAM));
        }
    }

}
